from galvatron.core import GeesProfiler, initialize_galvatron
import os

if __name__ == '__main__':
    args = initialize_galvatron(mode='profile_hardware')
    print(args)
    profiler = GeesProfiler(args)
    path = os.path.dirname(os.path.abspath(__file__))
    profiler.set_path(path)

    # # 使用nccl-test
    # profiler.profile_bandwidth(backend="nccl")

    # 使用Pytorch Profiler
    profiler.profile_bandwidth(backend="torch")
    # profile overlapping slowdown coefficient
